{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import pickle\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from itertools import chain\n",
    "from sentence import Sentence\n",
    "from sentence import TagPrefix\n",
    "from sentence import TagSurfix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class DataHandler(object):\n",
    "    def __init__(self, rootDir=None, dict_path=None, train_data=None):\n",
    "        self.rootDir = rootDir\n",
    "        self.dict_path = dict_path\n",
    "        self.train_data = train_data\n",
    "        self.spiltChar = ['。', '!', '！', '？', '?']\n",
    "        self.max_len = 200\n",
    "        self.totalLine = 0\n",
    "        self.longLine = 0\n",
    "        self.totalChars = 0\n",
    "        self.TAGPRE = TagPrefix.convert()\n",
    "    \n",
    "    def loadData(self):\n",
    "        self.loadRawData()\n",
    "        self.handlerRawData()\n",
    "    \n",
    "    def loadRawData(self):\n",
    "        self.datas = list()\n",
    "        self.labels = list()\n",
    "        if self.rootDir:\n",
    "            print(self.rootDir)\n",
    "            for dirName, subdirList, fileList in os.walk(self.rootDir):\n",
    "                # curDir = os.path.join(self.rootDir, dirName)\n",
    "                for file in fileList:                   \n",
    "                    if file.endswith(\".txt\"): \n",
    "                        curFile = os.path.join(dirName, file)\n",
    "                        print(\"processing:%s\" % (curFile))\n",
    "                        with open(curFile, \"r\", encoding='utf-8') as fp:\n",
    "                            for line in fp.readlines():\n",
    "                                self.processLine(line)\n",
    "\n",
    "            print(\"total:%d, long lines:%d, total chars:%d\" % (self.totalLine, self.longLine, self.totalChars))\n",
    "            print('Length of datas is %d' % len(self.datas))\n",
    "            print('Example of datas: ', self.datas[0])\n",
    "            print('Example of labels:', self.labels[0])\n",
    "\n",
    "    def processLine(self, line):\n",
    "        line = line.strip()  \n",
    "        nn = len(line)\n",
    "        seeLeftB = False  \n",
    "        start = 0\n",
    "        sentence = Sentence()  \n",
    "        try:\n",
    "            for i in range(nn):  \n",
    "                if line[i] == ' ':\n",
    "                    if not seeLeftB:\n",
    "                        token = line[start:i]\n",
    "                        if token.startswith('['):\n",
    "                            token_ = ''\n",
    "                            for j in [i.split('/') for i in token.split('[')[1].split(']')[0].split(' ')]:\n",
    "                                token_ += j[0]\n",
    "                            token_ = token_ + '/' + token.split('/')[-1]\n",
    "                            self.processToken(token_, sentence, False)\n",
    "                        else:\n",
    "                            self.processToken(token, sentence, False)\n",
    "                        start = i + 1\n",
    "                elif line[i] == '[':\n",
    "                    seeLeftB = True\n",
    "                elif line[i] == ']':\n",
    "                    seeLeftB = False\n",
    "        # 此部分未与上面处理方式统一，（小概率事件）数据多元化，增加模型泛化能力。\n",
    "            if start < nn:\n",
    "                token = line[start:]\n",
    "                if token.startswith('['):\n",
    "                    tokenLen = len(token)\n",
    "                    while tokenLen > 0 and token[tokenLen - 1] != ']':\n",
    "                        tokenLen = tokenLen - 1\n",
    "                    token = token[1:tokenLen - 1]\n",
    "                    ss = token.split(' ')\n",
    "                    ns = len(ss)\n",
    "                    for i in range(ns - 1):\n",
    "                        self.processToken(ss[i], sentence, False)\n",
    "                    self.processToken(ss[-1], sentence, True)\n",
    "                else:\n",
    "                    self.processToken(token, sentence, True)\n",
    "        except Exception as e:\n",
    "            print('处理数据异常, 异常行为：' + line)\n",
    "            print(e)\n",
    "\n",
    "    def processToken(self, tokens, sentence, end):\n",
    "        nn = len(tokens)\n",
    "        while nn > 0 and tokens[nn - 1] != '/':\n",
    "            nn = nn - 1\n",
    "\n",
    "        token = tokens[:nn - 1].strip() \n",
    "        tagPre = tokens[nn:].strip()\n",
    "        tagPre = self.TAGPRE.get(tagPre, TagPrefix.general.value) \n",
    "        if token not in self.spiltChar:\n",
    "            sentence.addToken(token, tagPre)\n",
    "        if token in self.spiltChar or end:\n",
    "            if sentence.chars > self.max_len: \n",
    "                self.longLine += 1\n",
    "            else:\n",
    "                x = []\n",
    "                y = []\n",
    "                self.totalChars += sentence.chars\n",
    "                sentence.generate_tr_line(x, y)\n",
    "\n",
    "                if len(x) > 0 and len(x) == len(y):\n",
    "                    self.datas.append(x)\n",
    "                    self.labels.append(y)\n",
    "                else:\n",
    "                    print('处理一行数据异常, 异常行如下')\n",
    "                    print(sentence.tokens)\n",
    "            self.totalLine += 1\n",
    "            sentence.clear()\n",
    "\n",
    "    def handlerRawData(self):\n",
    "        self.df_data = pd.DataFrame({'words': self.datas, 'tags': self.labels}, index=range(len(self.datas)))\n",
    "        self.df_data['sentence_len'] = self.df_data['words'].apply(\n",
    "            lambda words: len(words)) \n",
    "\n",
    "        all_words = list(chain(*self.df_data['words'].values))\n",
    "        sr_allwords = pd.Series(all_words)\n",
    "        sr_allwords = sr_allwords.value_counts() \n",
    "\n",
    "        set_words = sr_allwords.index\n",
    "        set_ids = range(1, len(set_words) + 1) \n",
    "        tags = ['x']\n",
    "\n",
    "        for _, memberPre in TagPrefix.__members__.items():\n",
    "            for _, memberSuf in TagSurfix.__members__.items():\n",
    "                if memberSuf is TagSurfix.S and memberPre is TagPrefix.general:\n",
    "                    tags.append(memberPre.value + memberSuf.value)\n",
    "                elif memberSuf != TagSurfix.S:\n",
    "                    tags.append(memberPre.value + memberSuf.value)\n",
    "\n",
    "        tags = list(set(tags))\n",
    "        print(tags)\n",
    "\n",
    "        tag_ids = range(len(tags))\n",
    "\n",
    "        self.word2id = pd.Series(set_ids, index=set_words)\n",
    "        self.id2word = pd.Series(set_words, index=set_ids)\n",
    "        self.id2word[len(set_ids) + 1] = '<NEW>'\n",
    "        self.word2id['<NEW>'] = len(set_ids) + 1\n",
    "\n",
    "        self.tag2id = pd.Series(tag_ids, index=tags)\n",
    "        self.id2tag = pd.Series(tags, index=tag_ids)\n",
    "\n",
    "        self.df_data['X'] = self.df_data['words'].apply(self.X_padding)\n",
    "        self.df_data['y'] = self.df_data['tags'].apply(self.y_padding)\n",
    "\n",
    "        self.X = np.asarray(list(self.df_data['X'].values))\n",
    "        self.y = np.asarray(list(self.df_data['y'].values))\n",
    "        print('X.shape={}, y.shape={}'.format(self.X.shape, self.y.shape))\n",
    "        print('Example of words: ', self.df_data['words'].values[0])\n",
    "        print('Example of X: ', self.X[0])\n",
    "        print('Example of tags: ', self.df_data['tags'].values[0])\n",
    "        print('Example of y: ', self.y[0])\n",
    "\n",
    "        with open(self.dict_path, 'wb') as outp:\n",
    "            pickle.dump(self.word2id, outp)\n",
    "            pickle.dump(self.id2word, outp)\n",
    "            pickle.dump(self.tag2id, outp)\n",
    "            pickle.dump(self.id2tag, outp)\n",
    "        print('** Finished saving the dict.')\n",
    "        \n",
    "        with open(self.train_data, 'wb') as outp:\n",
    "            pickle.dump(self.X, outp)\n",
    "            pickle.dump(self.y, outp)\n",
    "        print('** Finished saving the train data.')\n",
    "\n",
    "    def X_padding(self, words):\n",
    "\n",
    "        ids = list(self.word2id[words])\n",
    "        if len(ids) >= self.max_len: \n",
    "            return ids[:self.max_len]\n",
    "        ids.extend([0] * (self.max_len - len(ids))) \n",
    "        return ids\n",
    "\n",
    "    def y_padding(self, tags):\n",
    "        ids = list(self.tag2id[tags])\n",
    "        if len(ids) >= self.max_len:  \n",
    "            return ids[:self.max_len]\n",
    "        ids.extend([0] * (self.max_len - len(ids))) \n",
    "        return ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "../corpus\n",
      "processing:../corpus\\c1002-23995935.txt\n",
      "processing:../corpus\\c1002-23996898.txt\n",
      "total:27, long lines:0, total chars:1289\n",
      "Length of datas is 27\n",
      "Example of datas:  ['人', '民', '网', '1', '月', '1', '日', '讯', '据', '《', '纽', '约', '时', '报', '》', '报', '道', '，', '美', '国', '华', '尔', '街', '股', '市', '在', '2', '0', '1', '3', '年', '的', '最', '后', '一', '天', '继', '续', '上', '涨', '，', '和', '全', '球', '股', '市', '一', '样', '，', '都', '以', '最', '高', '纪', '录', '或', '接', '近', '最', '高', '纪', '录', '结', '束', '本', '年', '的', '交', '易']\n",
      "Example of labels: ['b', 'm', 'e', 'Date_b', 'Date_m', 'Date_m', 'Date_e', 's', 's', 's', 'b', 'm', 'm', 'e', 's', 'b', 'e', 's', 'b', 'e', 'b', 'm', 'e', 'b', 'e', 's', 'Date_b', 'Date_m', 'Date_m', 'Date_m', 'Date_e', 's', 'b', 'e', 'b', 'e', 'b', 'e', 'b', 'e', 's', 's', 'b', 'm', 'm', 'e', 'b', 'e', 's', 's', 's', 'b', 'm', 'm', 'e', 's', 'b', 'e', 'b', 'm', 'm', 'e', 'b', 'e', 's', 's', 's', 'b', 'e']\n",
      "['e', 'Date_m', 's', 'Date_e', 'm', 'Date_b', 'b', 'x']\n",
      "X.shape=(27, 200), y.shape=(27, 200)\n",
      "Example of words:  ['人', '民', '网', '1', '月', '1', '日', '讯', '据', '《', '纽', '约', '时', '报', '》', '报', '道', '，', '美', '国', '华', '尔', '街', '股', '市', '在', '2', '0', '1', '3', '年', '的', '最', '后', '一', '天', '继', '续', '上', '涨', '，', '和', '全', '球', '股', '市', '一', '样', '，', '都', '以', '最', '高', '纪', '录', '或', '接', '近', '最', '高', '纪', '录', '结', '束', '本', '年', '的', '交', '易']\n",
      "Example of X:  [ 14  12 221   6  58   6  26 320 145  70 162 213 124  28  65  28  77   1\n",
      "  90   9 192 179 216  55  86  10  18  49   6  39   5   2  31 109   4 421\n",
      "  67  81  33  72   1   3  69 281  55  86   4 199   1 438  30  31  42 239\n",
      " 228 371 352 361  31  42 239 228  53 215 184   5   2 311 267   0   0   0\n",
      "   0   0   0   0   0   0   0   0   0   0   0   0   0   0   0   0   0   0\n",
      "   0   0   0   0   0   0   0   0   0   0   0   0   0   0   0   0   0   0\n",
      "   0   0   0   0   0   0   0   0   0   0   0   0   0   0   0   0   0   0\n",
      "   0   0   0   0   0   0   0   0   0   0   0   0   0   0   0   0   0   0\n",
      "   0   0   0   0   0   0   0   0   0   0   0   0   0   0   0   0   0   0\n",
      "   0   0   0   0   0   0   0   0   0   0   0   0   0   0   0   0   0   0\n",
      "   0   0   0   0   0   0   0   0   0   0   0   0   0   0   0   0   0   0\n",
      "   0   0]\n",
      "Example of tags:  ['b', 'm', 'e', 'Date_b', 'Date_m', 'Date_m', 'Date_e', 's', 's', 's', 'b', 'm', 'm', 'e', 's', 'b', 'e', 's', 'b', 'e', 'b', 'm', 'e', 'b', 'e', 's', 'Date_b', 'Date_m', 'Date_m', 'Date_m', 'Date_e', 's', 'b', 'e', 'b', 'e', 'b', 'e', 'b', 'e', 's', 's', 'b', 'm', 'm', 'e', 'b', 'e', 's', 's', 's', 'b', 'm', 'm', 'e', 's', 'b', 'e', 'b', 'm', 'm', 'e', 'b', 'e', 's', 's', 's', 'b', 'e']\n",
      "Example of y:  [6 4 0 5 1 1 3 2 2 2 6 4 4 0 2 6 0 2 6 0 6 4 0 6 0 2 5 1 1 1 3 2 6 0 6 0 6\n",
      " 0 6 0 2 2 6 4 4 0 6 0 2 2 2 6 4 4 0 2 6 0 6 4 4 0 6 0 2 2 2 6 0 0 0 0 0 0\n",
      " 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      " 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      " 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0\n",
      " 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "** Finished saving the dict.\n",
      "** Finished saving the train data.\n",
      "[[ 14  12 221 ...,   0   0   0]\n",
      " [ 70 162 213 ...,   0   0   0]\n",
      " [ 87   6  18 ...,   0   0   0]\n",
      " ..., \n",
      " [ 19  23  46 ...,   0   0   0]\n",
      " [ 10   7   8 ...,   0   0   0]\n",
      " [  7   8  24 ...,   0   0   0]]\n",
      "<class 'numpy.ndarray'>\n",
      "(27, 200)\n"
     ]
    }
   ],
   "source": [
    "data = DataHandler(rootDir='../corpus', dict_path='../data/your_dict.pkl', train_data='../data/your_train_data.pkl')\n",
    "data.loadData()\n",
    "\n",
    "print(data.X)\n",
    "print(type(data.X))\n",
    "print(data.X.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
